# 指定文件编码为UTF-8
# coding: utf-8

"""
weight_init_compare.py 功能说明：
1. 比较不同权重初始化方法对神经网络训练的影响
2. 测试三种初始化方法：标准差0.01、Xavier初始化和He初始化
3. 使用MNIST数据集和4层神经网络进行实验
4. 记录并可视化训练过程中的损失变化
5. 使用平滑曲线展示训练趋势

关键点：
- 权重初始化对深度神经网络训练至关重要
- 不同激活函数需要匹配对应的初始化方法
- Xavier初始化适合Sigmoid/Tanh激活函数
- He初始化适合ReLU激活函数
- 使用SGD优化器进行训练
"""

# 导入操作系统模块
import os
# 导入系统模块
import sys
# 添加父目录到系统路径，以便导入父目录中的模块
sys.path.append(os.pardir)

# 导入NumPy数值计算库
import numpy as np
# 导入matplotlib绘图库
import matplotlib.pyplot as plt
# 从dataset.mnist导入MNIST数据加载函数
from dataset.mnist import load_mnist
# 从common.util导入平滑曲线函数
from common.util import smooth_curve
# 从common.multi_layer_net导入多层网络类
from common.multi_layer_net import MultiLayerNet
# 从common.optimizer导入SGD优化器
from common.optimizer import SGD

# 0.加载MNIST数据 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 加载MNIST数据集(自动归一化)
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True)

# 设置训练参数
train_size = x_train.shape[0]  # 训练数据大小
batch_size = 128               # 批量大小
max_iterations = 2000          # 最大迭代次数

# 1.实验设置 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 定义三种权重初始化方法
weight_init_types = {
    'std=0.01': 0.01,      # 固定标准差0.01
    'Xavier': 'sigmoid',   # Xavier初始化(适合Sigmoid)
    'He': 'relu'           # He初始化(适合ReLU)
}

# 使用SGD优化器，学习率0.01
optimizer = SGD(lr=0.01)

# 初始化网络和损失记录
networks = {}      # 存储各初始化方法对应的网络
train_loss = {}    # 存储各初始化方法的损失记录

# 为每种初始化方法创建网络
for key, weight_type in weight_init_types.items():
    # 4层神经网络(784-100-100-100-100-10)
    networks[key] = MultiLayerNet(
        input_size=784,
        hidden_size_list=[100, 100, 100, 100],
        output_size=10,
        weight_init_std=weight_type  # 设置权重初始化方法
    )
    # 初始化损失记录列表
    train_loss[key] = []

# 2.训练过程 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
for i in range(max_iterations):
    # 随机选择批量数据
    batch_mask = np.random.choice(train_size, batch_size)
    x_batch = x_train[batch_mask]  # 批量输入数据
    t_batch = t_train[batch_mask]  # 批量标签数据

    # 对每种初始化方法进行参数更新
    for key in weight_init_types.keys():
        # 计算梯度
        grads = networks[key].gradient(x_batch, t_batch)
        # 更新参数
        optimizer.update(networks[key].params, grads)
        # 计算并记录损失
        loss = networks[key].loss(x_batch, t_batch)
        train_loss[key].append(loss)

    # 每100次迭代打印进度
    if i % 100 == 0:
        print("===========" + "iteration:" + str(i) + "===========")
        # 打印各初始化方法的当前损失
        for key in weight_init_types.keys():
            loss = networks[key].loss(x_batch, t_batch)
            print(key + ":" + str(loss))

# 3.结果可视化 ===# ... 省略其他代码 ...


# ==== 修改后代码 ====
# 设置各初始化方法的绘图标记样式
markers = {
    'std=0.01': 'o',  # 圆圈
    'Xavier': 's',    # 方形
    'He': 'D'         # 菱形
}

# 创建x轴数据(迭代次数)
x = np.arange(max_iterations)

# 绘制各初始化方法的损失曲线
for key in weight_init_types.keys():
    # 使用平滑曲线绘制，每100个点标记一次
    plt.plot(x, smooth_curve(train_loss[key]),
             marker=markers[key],
             markevery=100,
             label=key)

# 设置图表属性
plt.xlabel("iterations")  # x轴标签
plt.ylabel("loss")        # y轴标签
plt.ylim(0, 2.5)         # y轴范围
plt.legend()             # 显示图例
plt.show()              # 显示图表

"""
预期结果分析：
1. std=0.01:
   - 学习速度可能较慢
   - 可能陷入局部最优
2. Xavier:
   - 适合Sigmoid激活函数
   - 学习速度适中
3. He:
   - 适合ReLU激活函数
   - 通常收敛最快

注意事项：
1. 不同初始化方法适用于不同激活函数
2. 网络深度增加时，初始化方法的影响更明显
3. 可以尝试结合其他优化器(如Adam)观察效果
4. 实际应用中需要根据网络结构选择合适的初始化方法
"""
